import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from io import StringIO
from datetime import datetime, date


# 自定义函数用于数据预处理
def preprocess_data(data):
    # 去除空值
    data = data.dropna()
    # 处理异常值（假设异常值为负数）
    for col in ['confirmedIncr', 'curedIncr', 'deadIncr','suspectedCountIncr']:
        data = data[data[col] >= 0]
    # 将dateId列转换为datetime64[ns]类型
    data['dateId'] = pd.to_datetime(data['dateId'], format='%Y%m%d')
    # 按日期排序
    data = data.sort_values(by='dateId')
    return data


# 自定义函数用于计算总感染趋势
def calculate_total_trend(data):
    total_trend = data.groupby('dateId')['confirmedIncr'].sum().cumsum()
    return total_trend


# 自定义函数用于绘制总感染趋势图
def plot_total_trend(total_trend, start_date, end_date):
    filtered_trend = total_trend[(total_trend.index >= start_date) & (total_trend.index <= end_date)]
    plt.figure(figsize=(10, 6))
    plt.plot(filtered_trend.index, filtered_trend.values, label='Total Infected Trend')
    plt.title('Total COVID - 19 Infected Trend')
    plt.xlabel('Date')
    plt.ylabel('Total Infected Count')
    plt.legend(loc='upper left')
    plt.xticks(rotation=45)
    st.pyplot(plt)


# 自定义函数用于绘制各城市感染曲线图
def plot_city_trends(data, start_date, end_date, threshold):
    cities = data['id'].unique()
    plt.figure(figsize=(10, 6))
    for city in cities:
        city_data = data[data['id'] == city]
        filtered_data = city_data[(city_data['dateId'] >= start_date) & (city_data['dateId'] <= end_date)]
        plt.plot(filtered_data['dateId'], filtered_data['confirmedIncr'], label=city)
        # 标记高风险区域
        high_risk_data = filtered_data[filtered_data['confirmedIncr'] > threshold]
        if not high_risk_data.empty:
            plt.scatter(high_risk_data['dateId'], high_risk_data['confirmedIncr'], color='red',
                        label=f'{city} (High Risk)')
    plt.title('City - wise COVID - 19 Infected Trends')
    plt.xlabel('Date')
    plt.ylabel('New Infected Count')
    plt.legend(loc='upper left')
    plt.xticks(rotation=45)
    st.pyplot(plt)


# Streamlit应用
def main():
    st.title('COVID - 19 Data Analysis and Visualization')
    st.write('Upload a CSV file containing COVID - 19 data and perform analysis.')

    # 文件上传
    uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
    if uploaded_file is not None:
        try:
            # 读取数据并将dateId列转换为datetime64[ns]类型，注意这里读取字节数据并转换为字符串再解析为DataFrame
            file_content = uploaded_file.read().decode('utf-8')
            data = pd.read_csv(StringIO(file_content), parse_dates=['dateId'])
            # 数据预处理
            data = preprocess_data(data)

            # 时间范围选择
            min_date = data['dateId'].min()
            max_date = data['dateId'].max()
            start_date = st.date_input('Start Date', min_date)
            end_date = st.date_input('End Date', max_date)
            start_date = pd.to_datetime(start_date)
            end_date = pd.to_datetime(end_date)

            # 确保起始日期不晚于结束日期
            if start_date > end_date:
                st.error("起始日期不能晚于结束日期")
                return

            # 感染人数阈值输入
            threshold = st.number_input('Infection Threshold', min_value=0, value=100)

            # 计算总感染趋势
            total_trend = calculate_total_trend(data)

            # 绘制总感染趋势图
            st.subheader('Total COVID - 19 Infected Trend')
            plot_total_trend(total_trend, start_date, end_date)

            # 绘制各城市感染曲线图
            st.subheader('City - wise COVID - 19 Infected Trends')
            plot_city_trends(data, start_date, end_date, threshold)

            # 下载按钮，生成图像字节数据并传递正确类型的参数用于下载
            buf = StringIO()
            plt.savefig(buf, format='png')
            byte_data = buf.getvalue().encode('utf-8')
            st.download_button(
                label="Download Chart",
                data=byte_data,
                file_name="covid_trend.png",
                mime="image/png"
            )

        except Exception as e:
            st.error(f"Error processing file: {e}")


if __name__ == "__main__":
    main()